import os
import random

from cozepy import Coze, TokenAuth, COZE_CN_BASE_URL
from cozepy import Message, ChatStatus, MessageContentType, ChatEventType

coze_api_base = os.getenv("COZE_API_BASE") or COZE_CN_BASE_URL

# 通过环境变量引入密钥，访问 coze.cn 服务
coze = Coze(auth=TokenAuth("pat_9bSvwdF8Bgci5bec1dCVAISqhUt8VFezbiVCc6k01MSNMhaUMHl7A3KdYUAl9OvH"), base_url=coze_api_base)

# Create a bot instance in Coze, copy the last number from the web link as the bot's ID.
bot_id = "7427321297908301875"
# The user id identifies the identity of a user. Developers can use a custom business ID
# or a random string.
user_id = str(random.random())

# Call the coze.chat.stream method to create a chat. The create method is a streaming
# chat and will return a Chat Iterator. Developers should iterate the iterator to get
# chat event and handle them.
for event in coze.chat.stream(
    bot_id=bot_id,
    user_id=user_id,
    additional_messages=[
        Message.build_user_question_text("Tell a 500-word story."),
    ],
):
    if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
        print(event.message.content, end="", flush=True)

    if event.event == ChatEventType.CONVERSATION_CHAT_COMPLETED:
        print()
        print("token usage:", event.chat.usage.token_count)
